
import logging
import os
import sys
import torch
import pickle
import random

from torch.utils.data import TensorDataset
from tqdm import tqdm

logger = logging.getLogger(__name__)

class InputExample(object):
    """A single training/test example for simple sequence classification."""
    """创建一个输入实例
    Args:
        guid: 每个example拥有唯一的id
        text_a: 第一个句子的原始文本，一般对于文本分类来说，只需要text_a
        text_b: 第二个句子的原始文本，在句子对的任务中才有，分类问题中为None
        label: example对应的标签，对于训练集和验证集应非None，测试集为None
    """
    def __init__(self, guid, text, label=None):
        self.guid = guid
        self.text = text 
        self.label = label

class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, label_id, ori_tokens):

        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_id = label_id
        self.ori_tokens = ori_tokens


class NerProcessor(object):
    def sent2char(self,line):
        """
        句子处理成单词
        :param line: 原始行
        :return: 单词， 标签
        """
        res = line.strip('\n').split()
        return res
    def train_val_split(self,X, y, valid_size=0.2, random_state=2018, shuffle=True):
        """
        训练集验证集分割
        :param X: sentences
        :param y: labels
        :param random_state: 随机种子
        """
        logger.info('Train val split')

        data = []
        for data_x, data_y in tqdm(zip(X, y), desc='Merge'):
            data.append((data_x, data_y))
        del X, y

        N = len(data)
        test_size = int(N * valid_size)

        if shuffle:
            random.seed(random_state)
            random.shuffle(data)

        valid = data[:test_size]
        train = data[test_size:]
        return train, valid

    def read_data(self,input_file,labelpath):
        targets, sentences = [],[]
        with open(input_file, 'rb') as fr_1, open(labelpath, 'rb') as fr_2:
            sentences=pickle.load(fr_1)
            targets=pickle.load(fr_2)
            
        return sentences, targets
        #train=[(['有', '意', '为', '难'], ['O', 'O', 'O', 'O']),([],[])]
    def get_labels(self, args):
        labels = set()
        if os.path.exists(os.path.join(args.output_dir, "label_list.pkl")):
            logger.info(f"loading labels info from {args.output_dir}")
            with open(os.path.join(args.output_dir, "label_list.pkl"), "rb") as f:
                labels = pickle.load(f)
        else:
            # get labels from train data
            logger.info(f"loading labels info from train file and dump in {args.output_dir}")
            with open(args.train_label,'rb') as f:
                lines=pickle.load(f)
                for item in lines:
                    labels=labels|set(item.split())
                    
            if len(labels) > 0:
                with open(os.path.join(args.output_dir, "label_list.pkl"), "wb") as f:
                    pickle.dump(labels, f)
            else:
                logger.info("loading error and return the default labels B,I,O")
                labels = {"O", "B", "I"}
        
        return labels 

    def get_examples(self, input_file,labelpath):
        sent,targ = self.read_data(input_file,labelpath)
        examples = []

        for i, item in enumerate(zip(sent,targ)):
            if len(item[0].strip().split())<5:
                continue
            guid = str(i)
            text = item[0]#string
            label = item[1]

            examples.append(InputExample(guid=guid, text=text, label=label))
        #examples=[class1.text class1.label, class2, class3, ... , classn]
        return examples

def convert_examples_to_features(args, examples, label_list, max_seq_length, tokenizer):
    label_map = {label : i for i, label in enumerate(label_list)}#label2id
    #{'E-V': 0, 'B-N': 1, 'O': 2, 'E-N': 3, 'I-N': 4, 'B-V': 5}
    features = []

    for (ex_index, example) in tqdm(enumerate(examples), desc="convert examples"):#这是在实现给word打标签,一一对应label,子单词        
        textlist = example.text.split(" ")#['从', '这', '些', '反', '应', '中']
        labellist = example.label.split(" ")
        assert len(textlist) == len(labellist)
        if len(textlist)<5:
            continue
        tokens = []
        labels = []
        ori_tokens = []

        for i, word in enumerate(textlist):
            # 防止wordPiece情况出现，不过貌似不会
            token = tokenizer.tokenize(word)
            tokens.extend(token)
            label_1 = labellist[i]
            ori_tokens.append(word)

            # 单个字符不会出现wordPiece
            for m in range(len(token)):
                if m == 0:
                    labels.append(label_1)
                else:
                    if label_1 == "O":
                        labels.append("O")
                    else:#为 B 或者 I
                        labels.append("I")#如果有子单词，id加载为I
            
        if len(tokens) >= max_seq_length - 1:
            tokens = tokens[0:(max_seq_length - 2)]  # -2 的原因是因为序列需要加一个句首和句尾标志
            labels = labels[0:(max_seq_length - 2)]
            ori_tokens = ori_tokens[0:(max_seq_length - 2)]

        ori_tokens = ["[CLS]"] + ori_tokens + ["[SEP]"]
        
        ntokens = []
        segment_ids = []
        label_ids = []
        ntokens.append("[CLS]")
        segment_ids.append(0)
        label_ids.append(label_map["O"])#开始的CLS对应label指定为：O

        for i, token in enumerate(tokens):
            ntokens.append(token)
            segment_ids.append(0)
            label_ids.append(label_map[labels[i]])

        ntokens.append("[SEP]")
        segment_ids.append(0)
        label_ids.append(label_map["O"])#最后加的是CLS='O'
        input_ids = tokenizer.convert_tokens_to_ids(ntokens)   
        
        input_mask = [1] * len(input_ids)

        assert len(ori_tokens) == len(ntokens), f"{len(ori_tokens)}, {len(ntokens)}, {ori_tokens}"

        while len(input_ids) < max_seq_length:
            input_ids.append(0)
            input_mask.append(0)
            segment_ids.append(0)
            # we don't concerned about it!
            label_ids.append(0)#padding直接补零
            ntokens.append("**NULL**")

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length
        assert len(label_ids) == max_seq_length

        if ex_index < 5:
            logger.info("*** Example ***")
            logger.info("guid: %s" % (example.guid))
            logger.info("tokens: %s" % " ".join(
                [str(x) for x in ntokens]))
            logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
            logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
            logger.info("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
            logger.info("label_ids: %s" % " ".join([str(x) for x in label_ids]))

        features.append(
                InputFeatures(input_ids=input_ids,
                              input_mask=input_mask,
                              segment_ids=segment_ids,
                              label_id=label_ids,
                              ori_tokens=ori_tokens))#原始句子token,以SEP结尾,没有加padding

    return features

def get_Dataset(args, processor, tokenizer, mode="train"):
    if mode == "train":
        filepath = args.train_file
        labelpath = args.train_label
    elif mode == "eval":
        filepath = args.eval_file
        labelpath = args.eval_label
    elif mode == "test":
        filepath = args.test_file
        labelpath = args.test_label
    else:
        raise ValueError("mode must be one of train, eval, or test")

    examples = processor.get_examples(filepath,labelpath)#[class1[.text,.label],class2...classn]
    label_list = args.label_list

    features = convert_examples_to_features(
        args, examples, label_list, args.max_seq_length, tokenizer
    )

    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.long)

    data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)

    return examples, features, data



